import sqlite3

class KnowledgeBase:
    def __init__(self, db_path='knowledge_base.db'):
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()
        self._create_table()

    def _create_table(self):
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS knowledge_base (
            id INTEGER PRIMARY KEY,
            question TEXT NOT NULL,
            answer TEXT NOT NULL,
            category TEXT,
            source TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')
        self.conn.commit()

    def insert(self, question, answer, category="通用", source="内部整理"):
        self.cursor.execute(
            "INSERT INTO knowledge_base (question, answer, category, source) VALUES (?, ?, ?, ?)",
            (question, answer, category, source)
        )
        self.conn.commit()

    def search(self, question, top_n=3):
        query = """
        SELECT question, answer 
        FROM knowledge_base 
        WHERE question LIKE ? 
        ORDER BY created_at DESC 
        LIMIT ?
        """
        self.cursor.execute(query, (f"%{question}%", top_n))
        return self.cursor.fetchall()

    def close(self):
        self.conn.close()